import numpy as np
import math
import matplotlib.pyplot as plt

from bresham import draw_line_to_boundary, bresenham
from build_global_map import build_global_map
from get_sensor_init_track import gen_sensor_data


# 逆传感器模型更新函数（全局地图：仅更新障碍物）
def inverse_sensor_model_global(global_map, sensor_data_frames, robot_pose):
    # 获取机器人位姿
    robot_x, robot_y, robot_theta = robot_pose
    robot_x = robot_x * 10
    robot_y = robot_y * 10
    sensor_angles  = robot_theta
    # 遍历传感器数据
    for sensor_data in sensor_data_frames:
        # for i, distance in enumerate(sensor_data):
        # for data  in sensor_data:
            # angle = sensor_angles + robot_theta  # 传感器方向角
            
            # 计算障碍物的位置（终点栅格）
            # obstacle_x = int(robot_x + distance / RESOLUTION * np.cos(angle))
            # obstacle_y = int(robot_y + distance / RESOLUTION * np.sin(angle))
        obstacle_x = sensor_data[0]
        obstacle_y = sensor_data[1]
            
        # 更新障碍物栅格为占据概率
        if 0 <= obstacle_x < GLOBAL_MAP_SIZE[0] and 0 <= obstacle_y < GLOBAL_MAP_SIZE[1]:
            global_map[obstacle_x, obstacle_y] = P_OCCUPIED
    
    return global_map

# 逆传感器模型更新函数（局部地图：更新障碍物和射线）
def inverse_sensor_model_local(local_map, sensor_data_frames,  robot_pose):
    # 获取机器人位姿
    robot_x, robot_y, robot_theta = robot_pose
    robot_x = robot_x 
    robot_y = robot_y 
    sensor_angles = robot_theta
    
    # 局部地图的中心（机器人在局部地图中的位置）
    local_center_x = LOCAL_MAP_SIZE[0] // 2
    local_center_y = LOCAL_MAP_SIZE[1] // 2
    
    # 遍历传感器数据, 数据实在相对位置
    for sensor_data in sensor_data_frames:
        # angle = sensor_angles + robot_theta  # 传感器方向角、

        delta_x = sensor_data[0] - robot_x
        delta_y = sensor_data[1] - robot_y

        obstacle_x = int(local_center_x + delta_x)
        obstacle_y = int(local_center_y + delta_y)
        print(f'obts xy {obstacle_x, obstacle_y}')
        # 计算障碍物的位置（终点栅格，相对于局部地图中心）
        # obstacle_x = sensor_data[0]
        # obstacle_y = sensor_data[1]
        point1 = np.array([obstacle_x, obstacle_y])  # 点 A
        point2 = np.array([local_center_x, local_center_y])  # 点 B

        # 计算欧几里得距离
        distance = np.linalg.norm(point2 - point1)
        
        # 更新障碍物栅格为占据概率
        if 0 <= obstacle_x < LOCAL_MAP_SIZE[0] and 0 <= obstacle_y < LOCAL_MAP_SIZE[1]:
            local_map[obstacle_x, obstacle_y] = P_OCCUPIED
        
        # 更新空闲区域栅格为空闲概率（射线）
        # steps = int(distance / RESOLUTION)  # 从机器人到障碍物的栅格数量
        # for step in range(steps):
        #     free_x = int(local_center_x + step * np.cos(angle))
        #     free_y = int(local_center_y + step * np.sin(angle))

        # find_points = draw_line_to_boundary()
        free_points = bresenham(local_center_x, local_center_y, obstacle_x, obstacle_y)
        for points in free_points:
            free_x, free_y = points
            if 0 <= free_x < LOCAL_MAP_SIZE[0] and 0 <= free_y < LOCAL_MAP_SIZE[1]:
                local_map[free_x, free_y] = P_FREE
    
    return local_map

# 提取局部地图
def extract_local_map(global_map, robot_pose, local_map_size):
    robot_x, robot_y, _ = robot_pose
    robot_x = robot_x 
    robot_y = robot_y 
    half_size = local_map_size[0] // 2  # 局部地图半径

    # 计算局部地图范围，并将索引转换为整数
    x_min = max(int(robot_x - half_size), 0)
    x_max = min(int(robot_x + half_size), GLOBAL_MAP_SIZE[0] // 2)
    y_min = max(int(robot_y - half_size), 0)
    y_max = min(int(robot_y + half_size), GLOBAL_MAP_SIZE[1] // 2)

    # 提取局部地图
    local_map = global_map[x_min:x_max, y_min:y_max]

    # 如果局部地图不够大小，进行填充
    padded_map = np.full(local_map_size, 0.5)  # 填充为未知区域
    padded_map[:local_map.shape[0], :local_map.shape[1]] = local_map

    return padded_map

# 可视化地图并保存
def visualize_and_save_map(map_data, save_path, frame_number, map_type="Global"):
    plt.figure(figsize=(10, 10))
    plt.imshow(map_data, cmap='gray', origin='lower')
    plt.colorbar(label='Occupancy Probability')
    plt.title(f"{map_type} Map - Frame {frame_number}")
    
    # 保存地图到文件
    plt.savefig(save_path, dpi=300, bbox_inches='tight')  # 保存为高分辨率图片
    print(f"{map_type} Map saved to {save_path}")
    
    # 关闭绘图以释放内存
    plt.close()









# 主程序
if __name__ == '__main__':

    # 定义地图大小和分辨率
    GLOBAL_MAP_SIZE = (200, 200)  # 全局地图大小 (200x200 栅格)
    LOCAL_MAP_SIZE = (100, 100)  # 局部地图大小 (100x100 栅格)
    RESOLUTION = 0.1             # 每个栅格的边长为 0.1 米

    real_map = build_global_map(GLOBAL_MAP_SIZE, RESOLUTION)
    global_map = np.full(GLOBAL_MAP_SIZE, 0.5)
    local_map = np.full(LOCAL_MAP_SIZE, 0.5)

    # 定义传感器参数
    P_OCCUPIED = 0.9  # 栅格被占据的概率
    P_FREE = 0.1      # 栅格为空闲的概率

    robot_init_pose = [0, 100, 0]  # (x, y, 方向角)，机器人位于全局地图中心，朝向角度为 0
    robot_speed = 0.1  # 机器人移动速度 (单位：米/秒)
    sensor_data_frames, robot_pose_frames = gen_sensor_data(real_map, robot_init_pose, robot_speed)



    # 更新地图并可视化所有帧数据
    for frame_number, (sensor_data, robot_pose) in enumerate(zip(sensor_data_frames, robot_pose_frames), start=1):
        # 更新机器人位置（直线运动）
        # robot_pose[0] += robot_speed / RESOLUTION * np.cos(robot_pose[2])  # 更新 x 坐标
        # robot_pose[1] += robot_speed / RESOLUTION * np.sin(robot_pose[2])  # 更新 y 坐标
        
        
        # 使用逆传感器模型更新全局地图（仅障碍物）
        global_map = inverse_sensor_model_global(global_map, sensor_data,  robot_pose)
        
        # 提取局部地图
        local_map = extract_local_map(global_map, robot_pose, LOCAL_MAP_SIZE)
        
        # 使用逆传感器模型更新局部地图（障碍物和射线）
        local_map = inverse_sensor_model_local(local_map, sensor_data,  robot_pose)
        
        # 保存路径
        global_save_path = f"global/global_map_frame_{frame_number}.png"
        local_save_path = f"global/local_map_frame_{frame_number}.png"
        
        # 可视化并保存地图
        visualize_and_save_map(global_map, global_save_path, frame_number, map_type="Global")
        visualize_and_save_map(local_map, local_save_path, frame_number, map_type="Local")
        break
